import unittest
from quantsim import *
from math import pi, sqrt

class MockHeap(object):
    def __init__(self):
        self.calls = []
        
    def __register_call(self, method_name, arguments):
        call = (method_name, arguments)
        self.calls.append(call)
        
    def nnot(self, index):
        arguments = (index,)
        self.__register_call("nnot", arguments)
        
    def cnot(self, target_index, test_index):
        arguments = (target_index, test_index)
        self.__register_call("cnot", arguments)
        
    def cphase(self, a, indices, reverse=False):
        arguments = (a, indices, reverse)
        self.__register_call("cphase", arguments)
        
    def mix(self, index):
        arguments = (index,)
        self.__register_call("mix", arguments)
        
    def measure(self, index):
        arguments = (index,)
        self.__register_call("measure", arguments)
        
    def substate(self, index, size):
        return {1:1}

class TestQHeap(unittest.TestCase):
        
    def setUp(self):
        heap_size = 5
        self.heap_size = heap_size
        self.heap = QHeap(heap_size)
        self.initial_state_str = str(complex(1, 0)) + " |00000>"

    def test_QHeap_must_be_initialized_with_nonnegative_number_of_qubits(self):
        self.assertRaises(ValueError, QHeap, -1)

    def test_QHeap_must_be_initialized_with_nonzero_number_of_qubits(self):
        self.assertRaises(ValueError, QHeap, 0)
        
    def test_QHeap_initializes_to_correct_quantum_state(self):
        heap = self.heap
        self.assertEqual(str(heap), self.initial_state_str, "QHeap does not initialize to all zeros.")
        
    def test_nnot_raises_exception_if_index_is_too_high(self):
        heap_size = self.heap_size
        heap = self.heap
        self.assertRaises(IndexError, heap.nnot, heap_size)

    def test_nnot_raises_exception_if_index_is_negative(self):
        heap = self.heap
        self.assertRaises(IndexError, heap.nnot, -1)

    def test_nnot_inverts_correct_qubit(self):
        heap = self.heap
        heap.nnot(0)
        self.assertEqual(str(heap), str(complex(1, 0)) + " |00001>", "nnot did not invert the correct qubit.")

    def test_nnot_is_its_own_inverse(self):
        heap = self.heap
        heap.nnot(0)
        heap.nnot(0)
        self.assertEqual(str(heap), self.initial_state_str, "applying nnot twice did not return heap to original state")

    def test_cnot_raises_exception_if_either_index_is_out_of_range(self):
        heap = self.heap
        heap_size = self.heap_size
        self.assertRaises(IndexError, heap.cnot, heap_size, [0])
        self.assertRaises(IndexError, heap.cnot, 0, [heap_size])
        self.assertRaises(IndexError, heap.cnot, heap_size, [0, -1])
        self.assertRaises(IndexError, heap.cnot, -1, [0, heap_size])

    def test_cnot_is_its_own_inverse(self):
        heap = self.heap
        heap.cnot(1, [0])
        heap.cnot(1, [0])
        self.assertEqual(str(heap), self.initial_state_str, "applying cnot twice did not return heap to original state")

    def test_cnot_inverts_target_qubit_if_all_test_qubits_are_1(self):
        heap = self.heap
        heap.nnot(0)
        heap.nnot(1)
        heap.cnot(2, [0, 1])
        self.assertEqual(str(heap), str(complex(1, 0)) + " |00111>", "cnot did not invert the correct qubit")

    def test_cnot_ignores_target_qubit_if_any_test_qubits_are_0(self):
        heap = self.heap
        heap.nnot(0)
        heap.cnot(2, [0, 1])
        self.assertEqual(str(heap), str(complex(1, 0)) + " |00001>", "cnot inverted the incorrect qubit")

    def test_cphase_raises_exception_if_any_index_is_too_high_or_is_negative(self):
        heap = self.heap
        heap_size = self.heap_size
        self.assertRaises(IndexError, heap.cphase, 0, [0, heap_size])
        self.assertRaises(IndexError, heap.cphase, 0, [0, -1])

    def test_cphase_rotates_amplitude_if_all_test_qubits_are_1(self):
        heap = self.heap
        heap.nnot(0)
        heap.nnot(1)
        heap.cphase(pi / 4, [0, 1])
        heap_states = heap.states
        amplitude = heap_states[3]
        self.assertAlmostEqual(amplitude.real, sqrt(0.5), msg="cphase did not rotate the amplitude")
        self.assertAlmostEqual(amplitude.imag, sqrt(0.5), msg="cphase did not rotate the amplitude")

    def test_cphase_does_not_rotate_amplitude_if_any_test_qubits_are_0(self):
        heap = self.heap
        heap.nnot(0)
        heap.cphase(pi / 4, [0, 1])
        heap_states = heap.states
        amplitude = heap_states[1]
        self.assertAlmostEqual(amplitude.real, 1, msg="cphase rotated the amplitude when test qubit was 0")
        self.assertAlmostEqual(amplitude.imag, 0, msg="cphase rotated the amplitude when test qubit was 0")

    def test_cphase_rotates_in_opposite_direction_when_reverse_flag_is_set(self):
        heap = self.heap
        heap.nnot(0)
        heap.cphase(pi / 4, [0], reverse = True)
        heap_states = heap.states
        amplitude = heap_states[1]
        self.assertAlmostEqual(amplitude.real, sqrt(0.5), msg="cphase did not rotate the amplitude in reverse")
        self.assertAlmostEqual(amplitude.imag, -sqrt(0.5), msg="cphase did not rotate the amplitude in reverse")

    def test_cphase_and_cphase_reverse_will_cancel_each_other_out(self):
        heap = self.heap
        heap.nnot(0)
        heap.cphase(pi / 4, [0])
        heap.cphase(pi / 4, [0], reverse = True)
        heap_states = heap.states
        amplitude = heap_states[1]
        self.assertAlmostEqual(amplitude.real, 1, msg="cphase and cphase reverse did not invert each other")
        self.assertAlmostEqual(amplitude.imag, 0, msg="cphase and cphase reverse did not invert each other")

    def test_mix_raises_exception_if_index_is_too_high_or_is_negative(self):
        heap = self.heap
        heap_size = self.heap_size
        self.assertRaises(IndexError, heap.mix, heap_size)
        self.assertRaises(IndexError, heap.mix, -1)

    def test_mix_divides_initial_state_correctly(self):
        heap = self.heap
        heap.mix(0)
        heap_states = heap.states
        amplitude1 = heap_states[0]
        amplitude2 = heap_states[1]
        self.assertAlmostEqual(amplitude1.real, sqrt(0.5))
        self.assertAlmostEqual(amplitude2.real, sqrt(0.5))

    def test_mix_divides_state_with_target_qubit_equal_to_1_correctly(self):
        heap = self.heap
        heap.nnot(0)
        heap.mix(0)
        heap_states = heap.states
        amplitude1 = heap_states[0]
        amplitude2 = heap_states[1]
        self.assertAlmostEqual(amplitude1.real, sqrt(0.5))
        self.assertAlmostEqual(amplitude2.real, -sqrt(0.5))

    def test_mix_is_its_own_inverse(self):
        heap = self.heap
        heap.mix(0)
        heap.mix(0)
        heap_states = heap.states
        self.assertAlmostEqual(heap_states[0].real, 1)
        self.assertAlmostEqual(heap_states[0].imag, 0)

    def test_reset_returns_heap_to_initial_state(self):
        heap = self.heap
        heap.mix(0)
        heap.nnot(1)
        heap.cphase(pi/3, [1])
        heap.mix(2)
        heap.reset()
        self.assertEqual(str(heap), self.initial_state_str, "Reset does not return heap to intial state.")

    def test_measure_raises_exception_if_index_is_too_high_or_too_low(self):
        heap = self.heap
        heap_size = self.heap_size
        self.assertRaises(IndexError, heap.measure, heap_size)
        self.assertRaises(IndexError, heap.measure, -1)


    def test_measure_independent_qubits(self):
        heap = self.heap
        heap.nnot(0)
        heap.mix(0)
        heap.nnot(1)
        heap.mix(1)
        heap.measure(0)
        states = heap.states
        num_nonzero_states = 0
        for amplitude in states.itervalues():
            if abs(amplitude) >= 1e-6:
                num_nonzero_states += 1
        self.assertEqual(num_nonzero_states, 2, "Measure did not collapse the wave function into the correct number of eigenstates")
        vector_total = 0
        for value in states.itervalues():
                vector_total += abs(value) ** 2
        self.assertAlmostEqual(vector_total, 1.0, msg = "Sum of vector magnitudes should be 1")

    def test_measure_entangled_states(self):
        heap = self.heap
        heap.mix(0)
        heap.cnot(1,[0])
        heap.measure(0)
        states = heap.states
        num_nonzero_states = 0
        for amplitude in states.itervalues():
            if abs(amplitude) >= 1e-6:
                num_nonzero_states += 1
        self.assertEqual(num_nonzero_states, 1, "Measure did not collapse the wave function into the correct number of eigenstates")
        vector_total = 0
        for value in states.itervalues():
                vector_total += abs(value) ** 2
        self.assertAlmostEqual(vector_total, 1.0, msg = "Vector magnitude of sole eigenstate does not equal 1")

    def test_substate_raises_error_if_index_is_negative(self):
        heap = self.heap
        negative_index = -1
        size = 3
        self.assertRaises(IndexError, heap.substate, negative_index, size)

    def test_substate_raises_error_if_size_is_nonpositive(self):
        heap = self.heap
        negative_size = -1
        zero_size = 0
        index = 0
        self.assertRaises(IndexError, heap.substate, index, negative_size)
        self.assertRaises(IndexError, heap.substate, index, zero_size)

    def test_substate_raises_error_if_size_plus_index_is_too_large(self):
        heap = self.heap
        heap_size = self.heap_size
        index = 1
        size = heap_size - index
        self.assertRaises(IndexError, heap.substate, index, size+1)

    def test_substate_has_correct_number_of_states(self):
        heap = self.heap
        index = 0
        size = 2
        substate = heap.substate(index, size)
        self.assertEqual(len(substate), 2**size, "Substate has an incorrect number of states")

    def test_substate_returns_correct_amplitudes(self):
        heap = self.heap
        heap.mix(0)
        heap.mix(1)
        heap.mix(2)
        substate = heap.substate(1,2)
        for i in range(4):
            self.assertEqual(abs(substate[i]), 0.5)

    def test_allocate_raises_error_if_size_is_too_big_or_small(self):
        heap = self.heap
        zero_size = 0
        negative_size = -1
        too_big_size = self.heap_size + 1
        self.assertRaises(IndexError, heap.allocate, zero_size)
        self.assertRaises(IndexError, heap.allocate, negative_size)
        self.assertRaises(MemoryError, heap.allocate, too_big_size)
         
    def test_allocate_returns_reg_of_correct_size(self):
        heap = self.heap
        reg_size = 2
        reg = heap.allocate(reg_size)
        self.assertEqual(len(reg), reg_size)

    def test_allocate_avoids_overlap(self):
        heap = self.heap
        reg_size = 2
        reg1 = heap.allocate(reg_size)
        reg2 = heap.allocate(reg_size)
        self.assert_(reg1[0].index + reg_size <= reg2[0].index)

    def test_allocate_raises_memoryerror_when_out_of_heap_space(self):
        heap = self.heap
        reg_size = self.heap_size
        heap.allocate(reg_size)
        self.assertRaises(MemoryError, heap.allocate, 1)
        
class TestQubit(unittest.TestCase):
    
    def test_qubit_has_index_and_heap_and_size(self):
        heap = QHeap(1)
        index = 0
        qubit = Qubit(index, heap)
        self.assertTrue(qubit.heap is heap, "Qubit.heap does not point to the correct heap")
        self.assertEqual(qubit.index, index, "Qubit.index points to the wrong index")
        self.assertEqual(len(qubit), 1, "Qubits must have size 1")
        
class TestQReg(unittest.TestCase):
    
    def setUp(self):
        self.heap = QHeap(3)
        self.index = 0
        self.size = 1
        self.reg = QReg(self.heap, self.index, self.size)
    
    def test_qreg_initializes_to_correct_size_index_and_heap(self):
        heap = self.heap
        index = self.index
        size = self.size
        reg = self.reg
        self.assertTrue(reg.heap is heap, "QReg does not initialize to the correct heap")
        self.assertEqual(reg.index, index, "QReg does not initialize to the correct index")
        self.assertEqual(len(reg), size, "QReg does not initialize to the correct size")
        
    def test_qreg_prints_correctly(self):
        heap = self.heap
        reg_str = str(self.reg)
        self.assertEqual(reg_str, "(1+0j) |0>")
        
class TestQuantumFunctions(unittest.TestCase):
    
    def setUp(self):
        mock_heap = MockHeap()
        index = 0
        size = 3
        reg = QReg(mock_heap, index, size)
        self.mock_heap = mock_heap
        self.reg = reg
        self.size = size
        self.index = 0
        
    def test_Not_calls_the_correct_methods(self):
        self.verify_function_with_single_index_as_argument_for_register(Not, "nnot")
        
    def test_Not_works_for_qubits(self):
        self.verify_function_with_single_index_as_argument_for_qubit(Not, "nnot")
        
    def test_CNot_raises_ValueError_if_arguments_are_from_different_heaps(self):
        mock_heap2 = MockHeap()
        reg_from_different_heap = QReg(mock_heap2, 0, 2)
        self.assertRaises(ValueError, CNot, self.reg[0], reg_from_different_heap[0])    
        
    def test_CNot_call_the_correct_methods(self):    
        mock_heap = self.mock_heap
        initial_test_index = self.index
        test_reg = self.reg
        initial_target_index = 4
        target_reg_size = 3
        target_reg = QReg(mock_heap, initial_target_index, target_reg_size)
        CNot(target_reg, test_reg)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), len(target_reg))
        
        expected_method_name = "cnot"
        target_indices = []
        for call in calls:
            method_name = call[0]
            self.assertEqual(expected_method_name, method_name)
            arguments = call[1]
            target_index = arguments[0]
            test_indices = arguments[1]
            test_indices.sort()
            self.assertEqual(test_indices, [0,1,2])
            target_indices.append(target_index)
        target_indices.sort()
        self.assertEqual(target_indices, [4, 5, 6])
        
    def test_CNot_works_for_qubits(self):
        mock_heap = self.mock_heap
        index = 0
        target_qubit = Qubit(index, mock_heap)
        reg = self.reg
        test_qubit = reg[0]
        CNot(target_qubit, test_qubit)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), 1)
        
        sole_call = calls[0]
        expected_method_name = "cnot"
        called_method_name = sole_call[0]
        self.assertEqual(expected_method_name, called_method_name)
        
        expected_arguments = (target_qubit.index, [test_qubit.index])
        called_arguments = sole_call[1]
        self.assertEqual(expected_arguments, called_arguments)
        
    def test_CPhase_calls_the_correct_methods(self):
        mock_heap = self.mock_heap
        initial_index = self.index
        reg = self.reg
        shift_angle = pi / 2
        CPhase(shift_angle, reg)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), 1)
        
        sole_call = calls[0]
        method_name = sole_call[0]
        arguments = sole_call[1]
        expected_method_name = "cphase"
        expected_arguments = (shift_angle, [qubit.index for qubit in reg], False)
        self.assertEqual(method_name, expected_method_name)
        self.assertEqual(arguments, expected_arguments)
        
    def test_CPhase_works_for_qubits(self):
        mock_heap = self.mock_heap
        reg = self.reg
        test_qubit = reg[0]
        shift_angle = pi / 2
        CPhase(shift_angle, test_qubit)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), 1)
        
        sole_call = calls[0]
        method_name = sole_call[0]
        arguments = sole_call[1]
        expected_method_name = "cphase"
        expected_arguments = (shift_angle, [test_qubit.index], False)
        self.assertEqual(method_name, expected_method_name)
        self.assertEqual(arguments, expected_arguments)
        
    def test_CPhase_works_for_reverse(self):
        mock_heap = self.mock_heap
        reg = self.reg
        test_qubit = reg[0]
        shift_angle = pi / 2
        CPhase(shift_angle, test_qubit, True)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), 1)
        
        sole_call = calls[0]
        method_name = sole_call[0]
        arguments = sole_call[1]
        expected_method_name = "cphase"
        expected_arguments = (shift_angle, [test_qubit.index], True)
        self.assertEqual(method_name, expected_method_name)
        self.assertEqual(arguments, expected_arguments)
        
    def test_Mix_calls_the_correct_methods(self):
        self.verify_function_with_single_index_as_argument_for_register(Mix, "mix")
        
    def test_Mix_works_for_a_single_qubit(self):
        self.verify_function_with_single_index_as_argument_for_qubit(Mix, "mix")
        
    def test_Measure_calls_the_correct_methods(self):
        self.verify_function_with_single_index_as_argument_for_register(Measure, "measure")
        
    def test_Measure_works_for_a_single_qubit(self):
        self.verify_function_with_single_index_as_argument_for_qubit(Measure, "measure")
        
    def test_Measure_returns_the_correct_state(self):
        reg = self.reg
        Not(reg[0])
        state = Measure(reg)
        self.assertEqual(state, 1)
        
        
    def verify_function_with_single_index_as_argument_for_register(self, func, expected_method):
        mock_heap = self.mock_heap
        reg = self.reg
        func(reg)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), len(reg))
        
        arguments = []
        for call in calls:
            method = call[0]
            self.assertEqual(method, expected_method)
            argument = call[1]
            arguments.append(argument)
        expected_arguments = [(qubit.index,) for qubit in reg]
        expected_arguments.sort()
        arguments.sort()       
        self.assertEqual(expected_arguments, arguments)
        
    def verify_function_with_single_index_as_argument_for_qubit(self, func, expected_method):
        mock_heap = self.mock_heap
        qubit = self.reg[0]
        func(qubit)
        
        calls = mock_heap.calls
        self.assertEqual(len(calls), 1)
        
        sole_call = calls[0]
        method = sole_call[0]
        self.assertEqual(method, expected_method)
        argument = sole_call[1]
        expected_argument = (qubit.index,)
        self.assertEqual(argument, expected_argument)
        
if __name__ == "__main__":
    unittest.main()
